package model

import (
	"time"

	"git.oschina.net/ant180/scan/config"
	"github.com/xlvector/dlog"
	"gopkg.in/mgo.v2"
)

type Mongodb struct {
	conf    *config.MongodbConf
	mclient *mgo.Session

	mongoPool  map[string]*IUPool
	poolHTFunc PoolHTFunc
}

type PoolHTFunc func([]interface{})

func NewMongodb(conf *config.MongodbConf) (*Mongodb, error) {
	m := &Mongodb{
		conf:       conf,
		mongoPool:  make(map[string]*IUPool),
		poolHTFunc: nil,
	}

	if err := m.connect(); err != nil {
		return nil, err
	}

	go m.poolHaertbeat()

	return m, nil
}

func NewMongodb2(conf *config.MongodbConf, poolhtfunc PoolHTFunc) (*Mongodb, error) {
	m := &Mongodb{
		conf:       conf,
		mongoPool:  make(map[string]*IUPool),
		poolHTFunc: poolhtfunc,
	}

	if err := m.connect(); err != nil {
		return nil, err
	}

	go m.poolHaertbeat()

	return m, nil
}

func (p *Mongodb) connect() error {
	info := &mgo.DialInfo{
		Addrs:    []string{p.conf.Addr},
		Timeout:  time.Duration(p.conf.Timeout) * time.Second,
		Database: p.conf.DB,
		Username: p.conf.User,
		Password: p.conf.PWD,
	}

	session, err := mgo.DialWithInfo(info)
	if err != nil {
		dlog.Error("mongodb dial fail! %v", err)
		return err
	}
	session.SetMode(mgo.Monotonic, true)

	p.mclient = session

	return nil
}

func (p *Mongodb) SaveData(collect string, data interface{}) error {
	db := p.conf.DB
	c := p.mclient.DB(db).C(collect)
	err := c.Insert(data)
	if err != nil {
		if err := p.connect(); err != nil {
			dlog.Error("mongodb connect fail: %v", err)
			return err
		}
		c = p.mclient.DB(db).C(collect)
		err = c.Insert(data)
		if err != nil {
			dlog.Error("insert fail: %v", err)
			return err
		}
	}

	return nil
}

func (p *Mongodb) SaveData2(db, collect string, data interface{}) error {
	c := p.mclient.DB(db).C(collect)
	err := c.Insert(data)
	if err != nil {
		if err := p.connect(); err != nil {
			dlog.Error("mongodb connect fail: %v", err)
			return err
		}
		c = p.mclient.DB(db).C(collect)
		err = c.Insert(data)
		if err != nil {
			dlog.Error("insert fail: %v", err)
			return err
		}
	}

	return nil
}

const (
	POOL_WAIT = 5
	POOL_NUM  = 20
)

func (p *Mongodb) SaveDataBatch(collect string, data []interface{}) error {
	db := p.conf.DB
	c := p.mclient.DB(db).C(collect)
	err := c.Insert(data...)
	if err != nil {
		if err := p.connect(); err != nil {
			dlog.Error("mongodb connect fail: %v", err)
			return err
		}
		c = p.mclient.DB(db).C(collect)
		err = c.Insert(data...)
		if err != nil {
			dlog.Error("insert fail: %v", err)
			return err
		}
	}

	return nil
}

func (p *Mongodb) SaveDataPool(collect string, data interface{}) error {

	colPool, ok := p.mongoPool[collect]
	if !ok {
		colPool = NewIUPool(POOL_WAIT, POOL_NUM)
		p.mongoPool[collect] = colPool
	}

	ok, err := colPool.Insert(data)
	if ok {
		datas := colPool.GetAll()
		return p.SaveDataBatch(collect, datas)
	}

	return err
}

func (p *Mongodb) poolHaertbeat() {
	tc := time.NewTicker(POOL_WAIT)
	for {
		<-tc.C
		for k, v := range p.mongoPool {
			if !v.IsDone() {
				continue
			}
			datas := v.GetAll()
			if len(datas) <= 0 {
				continue
			}
			err := p.SaveDataBatch(k, datas)
			dlog.Info("pool ht: %s %d", k, len(datas))
			if err != nil {
				dlog.Error("pool ht save fail! %v", err)
				if p.poolHTFunc != nil {
					p.poolHTFunc(datas)
				}
			}
		}
	}
}
